#!/usr/bin/python
# -*- encoding: utf8 -*-

class HeapError(Exception):
    pass

class Heap(object):
    '''
    A max heap with remove() and update() methods and their time complexities are O(log(N)).

    Note that remove() and update() fails when there are duplicated data.
    '''

    def __init__(self, comparator=cmp):
        self._cmp = comparator
        self._list = [None]
        self._index = {}

    def size(self):
        return len(self._list) - 1

    def push(self, obj):
        self._append(obj)
        self._heap_up(self._last_index())

    def pop(self):
        if not self.size():
            raise HeapError('pop from empty heap.')

        self._swap(1, self._last_index())
        max_obj = self._pop()
        self._heap_down(1)
        return max_obj

    def remove(self, obj):
        current = self._index[obj]
        self._swap(current, self._last_index())
        self._pop()

        self._heap_up(current)
        self._heap_down(current)

    def update(self, obj, new_obj):
        self.remove(obj)
        self.push(new_obj)

    def _heap_up(self, current):
        if current > self._last_index():
            return

        while current > 1:
            parent = current / 2
            if self._cmp(self._list[current], self._list[parent]) <= 0:
                break

            self._swap(current, parent)
            current = parent

    def _heap_down(self, current):
        while current * 2 <= self._last_index():
            max_child = current * 2
            child = max_child + 1
            if (child <= self._last_index()
                and self._cmp(self._list[child], self._list[max_child]) > 0):
                max_child = child

            if self._cmp(self._list[max_child], self._list[current]) <= 0:
                break

            self._swap(current, max_child)
            current = max_child

    def _append(self, obj):
        self._list.append(obj)
        self._index[obj] = self._last_index()

    def _pop(self):
        obj = self._list.pop()
        self._index.pop(obj)
        return obj

    def _swap(self, a, b):
        assert self._index[self._list[a]] == a
        assert self._index[self._list[b]] == b
        self._list[a], self._list[b] = self._list[b], self._list[a]
        self._index[self._list[a]] = a
        self._index[self._list[b]] = b

    def _last_index(self):
        return len(self._list) - 1
